#include <arpa/inet.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>

#include <iostream>
#include <set>

void get_min_max_ip_from_sockaddr(const struct sockaddr* sa, int mask, struct sockaddr* min_ip, struct sockaddr* max_ip)
{
    if (sa->sa_family == AF_INET) {
        struct sockaddr_in* sa_in = (struct sockaddr_in*)sa;
        struct sockaddr_in* min_ip_in = (struct sockaddr_in*)min_ip;
        struct sockaddr_in* max_ip_in = (struct sockaddr_in*)max_ip;

        uint32_t netmask = htonl((0xFFFFFFFF << (32 - mask)) & 0xFFFFFFFF);
        uint32_t network = sa_in->sin_addr.s_addr & netmask;
        uint32_t broadcast = network | ~netmask;

        memset(min_ip_in, 0, sizeof(struct sockaddr_in));
        min_ip_in->sin_family = AF_INET;
        min_ip_in->sin_addr.s_addr = network;

        memset(max_ip_in, 0, sizeof(struct sockaddr_in));
        max_ip_in->sin_family = AF_INET;
        max_ip_in->sin_addr.s_addr = broadcast;
    } else if (sa->sa_family == AF_INET6) {
        struct sockaddr_in6* sa_in6 = (struct sockaddr_in6*)sa;
        struct sockaddr_in6* min_ip_in6 = (struct sockaddr_in6*)min_ip;
        struct sockaddr_in6* max_ip_in6 = (struct sockaddr_in6*)max_ip;

        struct in6_addr netmask;
        memset(&netmask, 0, sizeof(netmask));
        for (int i = 0; i < mask / 8; ++i) {
            netmask.s6_addr[i] = 0xFF;
        }
        if (mask % 8 != 0) {
            netmask.s6_addr[mask / 8] = 0xFF << (8 - (mask % 8));
        }

        struct in6_addr network;
        for (int i = 0; i < 16; ++i) {
            network.s6_addr[i] = sa_in6->sin6_addr.s6_addr[i] & netmask.s6_addr[i];
        }

        struct in6_addr broadcast;
        for (int i = 0; i < 16; ++i) {
            broadcast.s6_addr[i] = network.s6_addr[i] | (~netmask.s6_addr[i]);
        }

        memset(min_ip_in6, 0, sizeof(struct sockaddr_in6));
        min_ip_in6->sin6_family = AF_INET6;
        min_ip_in6->sin6_addr = network;

        memset(max_ip_in6, 0, sizeof(struct sockaddr_in6));
        max_ip_in6->sin6_family = AF_INET6;
        max_ip_in6->sin6_addr = broadcast;
    } else {
        printf("%s:%d  family = xxxxxxxxxxx\n", __func__, __LINE__);
    }
}

typedef union {
    struct sockaddr_in in;
    struct sockaddr_in6 in6;
    struct sockaddr sa;
} cidr_sockaddr;

class CIDR {
public:
    CIDR(const std::string& addr, uint8_t mask)
    {
        cidr_sockaddr sa;
        set_sockaddr(sa, addr);
        get_min_max_ip_from_sockaddr(&sa.sa, mask, &start_.sa, &end_.sa);
        info_ = addr + "/" + std::to_string(mask);
        if (sa.sa.sa_family == AF_INET) {
            start_.in.sin_port = mask;
        } else {
            start_.in6.sin6_port = mask;
        }
    }

    bool operator<(const CIDR& rhs) const
    {
        if (start_.sa.sa_family == AF_INET) { // v4
            return memcmp(&start_.in.sin_addr.s_addr, &rhs.start_.in.sin_addr.s_addr, sizeof(in_addr)) < 0 && memcmp(&end_.in.sin_addr.s_addr, &rhs.start_.in.sin_addr.s_addr, sizeof(in_addr)) < 0;
        } else { // v6
            return memcmp(&start_.in6.sin6_addr.s6_addr, &rhs.start_.in6.sin6_addr.s6_addr, sizeof(in6_addr)) < 0 && memcmp(&end_.in6.sin6_addr.s6_addr, &rhs.start_.in6.sin6_addr.s6_addr, sizeof(in6_addr)) < 0;
        }
    }

    std::string get_start_ip() const
    {
        char ip_str[INET6_ADDRSTRLEN];
        memset(ip_str, '0', INET6_ADDRSTRLEN);
        if (start_.sa.sa_family == AF_INET) { // v4
            inet_ntop(AF_INET, &start_.in.sin_addr, ip_str, INET6_ADDRSTRLEN);
        } else {
            inet_ntop(AF_INET6, &start_.in6.sin6_addr, ip_str, INET6_ADDRSTRLEN);
        }
        return std::string(ip_str);
    }

    std::string get_end_ip() const
    {
        char ip_str[INET6_ADDRSTRLEN];
        memset(ip_str, '0', INET6_ADDRSTRLEN);
        if (end_.sa.sa_family == AF_INET) { // v4
            inet_ntop(AF_INET, &end_.in.sin_addr, ip_str, INET6_ADDRSTRLEN);
        } else {
            inet_ntop(AF_INET6, &end_.in6.sin6_addr, ip_str, INET6_ADDRSTRLEN);
        }
        return std::string(ip_str);
    }

    int get_prefix_length() const
    {
        if (start_.sa.sa_family == AF_INET) {
            return start_.in.sin_port;
        } else {
            return start_.in6.sin6_port;
        }
    }

    bool set_sockaddr(cidr_sockaddr& sa, const std::string& addr) const
    {
        bool bret = false;
        if (addr.find(':') != std::string::npos) { // IPv6
            if (1 == inet_pton(AF_INET6, addr.c_str(), &sa.in6.sin6_addr)) {
                sa.sa.sa_family = AF_INET6;
                bret = true;
            }
        } else { // IPv4
            if (1 == inet_pton(AF_INET, addr.c_str(), &sa.in.sin_addr)) {
                sa.sa.sa_family = AF_INET;
                bret = true;
            }
        }
        return bret;
    }

    bool contains(const std::string& ip) const
    {
        cidr_sockaddr tmp;
        if (!set_sockaddr(tmp, ip)) {
            return false;
        }
        if (start_.sa.sa_family != tmp.sa.sa_family) {
            return false;
        }
        uint8_t prefix_length = get_prefix_length();
        if (start_.sa.sa_family == AF_INET6) { // v6
            return memcmp(&start_.in6.sin6_addr, &tmp.in6.sin6_addr, sizeof(struct in6_addr)) <= 0
                && memcmp(&tmp.in6.sin6_addr, &end_.in6.sin6_addr, sizeof(struct in6_addr)) <= 0;

        } else {
            return memcmp(&start_.in.sin_addr, &tmp.in.sin_addr, sizeof(struct in_addr)) <= 0
                && memcmp(&tmp.in.sin_addr, &end_.in.sin_addr, sizeof(struct in_addr)) <= 0;
            uint32_t net = ntohl(start_.in.sin_addr.s_addr);
            uint32_t addr = ntohl(tmp.in.sin_addr.s_addr);
            uint32_t mask = ((uint32_t)(~0)) << (32 - prefix_length);
            return net == (addr & mask);
        }
    }

private:
    cidr_sockaddr start_;
    cidr_sockaddr end_;
    std::string info_;
};

void insert(std::set<CIDR>& cidr_set, const std::string& ip, int mask)
{
	CIDR tmp(ip,mask);
    cidr_set.lower_bound(CIDR(ip, mask));
    auto insert_result = cidr_set.insert(CIDR(ip, mask));
    printf("%s 插入%s/%d, set.size=%zd, start=%s,end=%s\n",
        insert_result.second ? "success" : "failed",
        ip.c_str(), mask,
        cidr_set.size(), tmp.get_start_ip().c_str(), tmp.get_end_ip().c_str());
}

void test_find(std::set<CIDR>& cidr_set, const std::string& ip, int mask)
{
    auto it = cidr_set.lower_bound(CIDR(ip, mask));
    if (it != cidr_set.end()) {
        if ((*it).contains(ip)) {
            printf("success find %s/%d in %s/%d\n", ip.c_str(), mask, (*it).get_start_ip().c_str(), (*it).get_prefix_length());
        } else {
            printf("failed find ip=%s, in=%s/%d\n", ip.c_str(), (*it).get_start_ip().c_str(), (*it).get_prefix_length());
        }
    } else {
        printf("not find %s/%d\n", ip.c_str(), mask);
    }
}

void print_size()
{
    printf("sizeof(sockaddr_storage)=%zd\n", sizeof(struct sockaddr_storage));
    printf("sizeof(sockaddr_in)=%zd\n", sizeof(struct sockaddr_in));
    printf("sizeof(sockaddr_in6)=%zd\n", sizeof(struct sockaddr_in6));
    printf("sizeof(sockaddr)=%zd\n", sizeof(struct sockaddr));
    printf("sizeof(CIDR)=%zd\n", sizeof(CIDR));
    printf("sizeof(unit8_t)=%zd\n", sizeof(uint8_t));
    printf("sizeof(sa_family_t)=%zd\n", sizeof(sa_family_t));
    printf("sizeof(in_port_t)=%zd\n", sizeof(in_port_t));
    CIDR c("192.168.0.0", 24);
    printf("c.mask=%d\n", c.get_prefix_length());
    printf("(~0)=%ld\n", (~0));
    printf("((uint32_t)(~0))=%u\n", ((uint32_t)(~0)));
}

int main()
{
    print_size();

    // 使用 CIDR 类的实例创建 std::set
    std::set<CIDR> cidr_set;

    insert(cidr_set, "10.4.136.0", 23);
    insert(cidr_set, "192.168.0.0", 24);
    insert(cidr_set, "192.168.0.128", 25);
    insert(cidr_set, "10.168.0.128", 25);
    insert(cidr_set, "1.1.0.0", 8);
    insert(cidr_set, "1.1.3.4", 8);
    insert(cidr_set, "2.2.2.2", 12);
    insert(cidr_set, "255.255.255.255", 32);

    insert(cidr_set, "2604:a880:8000::", 36);
    insert(cidr_set, "2604:a880:8000::", 36);
    insert(cidr_set, "2604:a880:8fff::", 36);
    insert(cidr_set, "2001:0db8:85a3::", 64);

    test_find(cidr_set, "255.255.255.255", 32);
    test_find(cidr_set, "192.168.0.128", 32);
    test_find(cidr_set, "192.168.0.255", 32);
    test_find(cidr_set, "192.168.1.255", 32);
    test_find(cidr_set, "2.2.2.4", 32);

    test_find(cidr_set, "2604:a880:8000::", 36);
    test_find(cidr_set, "2604:a880:8fff:ffff:ffff:ffff:ffaf:ffff", 128);
    test_find(cidr_set, "2604:a880:9::", 128);
    return 0;
}
